#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
脚本执行问题修复工具
用于诊断和修复项目中脚本的执行问题
"""

import os
import sys
import subprocess
import importlib.util
import json
from pathlib import Path
from datetime import datetime

class ExecutionIssueFixer:
    def __init__(self, project_root=None):
        self.project_root = Path(project_root) if project_root else Path(__file__).parent.parent
        self.issues_found = []
        self.fixes_applied = []
        
    def diagnose_issues(self):
        """诊断脚本执行问题"""
        print("🔍 开始诊断脚本执行问题...")
        
        # 检查Python环境
        self._check_python_environment()
        
        # 检查模块导入问题
        self._check_module_imports()
        
        # 检查文件权限
        self._check_file_permissions()
        
        # 检查依赖关系
        self._check_dependencies()
        
        return self.issues_found
    
    def _check_python_environment(self):
        """检查Python环境"""
        print("  📋 检查Python环境...")
        
        # 检查Python版本
        python_version = sys.version_info
        if python_version.major < 3 or (python_version.major == 3 and python_version.minor < 6):
            self.issues_found.append({
                "type": "python_version",
                "severity": "high",
                "description": f"Python版本过低: {python_version.major}.{python_version.minor}",
                "solution": "升级到Python 3.6或更高版本"
            })
        
        # 检查PYTHONPATH
        pythonpath = os.environ.get('PYTHONPATH', '')
        project_in_path = str(self.project_root) in pythonpath
        
        if not project_in_path:
            self.issues_found.append({
                "type": "pythonpath",
                "severity": "medium",
                "description": "项目根目录不在PYTHONPATH中",
                "solution": f"添加 {self.project_root} 到PYTHONPATH"
            })
    
    def _check_module_imports(self):
        """检查模块导入问题"""
        print("  📦 检查模块导入...")
        
        # 主要脚本文件
        scripts_to_check = [
            "apps/check.py",
            "apps/redteam_launcher.py", 
            "scripts/ultimate_persistence_tester.py",
            "modules/exploits/main_exploit.py"
        ]
        
        for script_path in scripts_to_check:
            full_path = self.project_root / script_path
            if full_path.exists():
                self._analyze_script_imports(full_path, script_path)
    
    def _analyze_script_imports(self, script_path, relative_path):
        """分析单个脚本的导入问题"""
        try:
            with open(script_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            # 查找相对导入
            import_lines = [line.strip() for line in content.split('\n') 
                          if line.strip().startswith(('from .', 'import .'))]
            
            if import_lines:
                self.issues_found.append({
                    "type": "relative_import",
                    "severity": "high", 
                    "description": f"{relative_path} 包含相对导入",
                    "details": import_lines,
                    "solution": "转换为绝对导入或添加__init__.py文件"
                })
            
            # 查找可能缺失的模块
            import_lines = [line.strip() for line in content.split('\n')
                          if line.strip().startswith(('import ', 'from '))]
            
            for line in import_lines:
                if 'modules.' in line or 'config.' in line or 'apps.' in line:
                    module_name = self._extract_module_name(line)
                    if module_name and not self._module_exists(module_name):
                        self.issues_found.append({
                            "type": "missing_module",
                            "severity": "high",
                            "description": f"模块 {module_name} 不存在或无法导入",
                            "file": relative_path,
                            "solution": "检查模块路径或创建缺失的模块文件"
                        })
                        
        except Exception as e:
            self.issues_found.append({
                "type": "file_read_error",
                "severity": "medium",
                "description": f"无法读取文件 {relative_path}: {str(e)}",
                "solution": "检查文件权限和编码"
            })
    
    def _extract_module_name(self, import_line):
        """从导入语句中提取模块名"""
        try:
            if import_line.startswith('from '):
                # from modules.xxx import yyy
                parts = import_line.split()
                if len(parts) >= 2:
                    return parts[1].split('.')[0]
            elif import_line.startswith('import '):
                # import modules.xxx
                parts = import_line.split()
                if len(parts) >= 2:
                    return parts[1].split('.')[0]
        except:
            pass
        return None
    
    def _module_exists(self, module_name):
        """检查模块是否存在"""
        try:
            # 检查是否为项目内模块
            module_path = self.project_root / module_name
            if module_path.exists() and module_path.is_dir():
                return True
            
            # 检查是否为标准库或已安装的包
            spec = importlib.util.find_spec(module_name)
            return spec is not None
        except:
            return False
    
    def _check_file_permissions(self):
        """检查文件权限"""
        print("  🔐 检查文件权限...")
        
        scripts_to_check = [
            "apps/check.py",
            "apps/redteam_launcher.py",
            "scripts/ultimate_persistence_tester.py"
        ]
        
        for script_path in scripts_to_check:
            full_path = self.project_root / script_path
            if full_path.exists():
                try:
                    # 在Windows上检查文件是否可读
                    with open(full_path, 'r') as f:
                        pass
                except PermissionError:
                    self.issues_found.append({
                        "type": "permission_error",
                        "severity": "high",
                        "description": f"文件 {script_path} 权限不足",
                        "solution": "修改文件权限或以管理员身份运行"
                    })
    
    def _check_dependencies(self):
        """检查依赖关系"""
        print("  📚 检查依赖关系...")
        
        # 常见的依赖模块
        common_deps = [
            'requests', 'psutil', 'colorama', 'cryptography',
            'paramiko', 'scapy', 'netifaces'
        ]
        
        missing_deps = []
        for dep in common_deps:
            try:
                importlib.import_module(dep)
            except ImportError:
                missing_deps.append(dep)
        
        if missing_deps:
            self.issues_found.append({
                "type": "missing_dependencies",
                "severity": "medium",
                "description": f"缺失依赖包: {', '.join(missing_deps)}",
                "solution": f"运行: pip install {' '.join(missing_deps)}"
            })
    
    def apply_fixes(self):
        """应用修复方案"""
        print("🔧 开始应用修复方案...")
        
        for issue in self.issues_found:
            if issue["type"] == "pythonpath":
                self._fix_pythonpath()
            elif issue["type"] == "relative_import":
                self._fix_relative_imports(issue)
            elif issue["type"] == "missing_dependencies":
                self._fix_dependencies(issue)
            elif issue["type"] == "missing_module":
                self._fix_missing_modules(issue)
    
    def _fix_pythonpath(self):
        """修复PYTHONPATH问题"""
        try:
            # 创建环境设置脚本
            env_script = self.project_root / "set_environment.py"
            with open(env_script, 'w', encoding='utf-8') as f:
                f.write(f"""#!/usr/bin/env python3
import sys
import os

# 添加项目根目录到Python路径
project_root = r"{self.project_root}"
if project_root not in sys.path:
    sys.path.insert(0, project_root)

# 设置环境变量
os.environ['PYTHONPATH'] = project_root + os.pathsep + os.environ.get('PYTHONPATH', '')

print(f"已设置项目环境: {{project_root}}")
""")
            
            self.fixes_applied.append({
                "type": "pythonpath",
                "description": "创建环境设置脚本",
                "file": str(env_script)
            })
            
        except Exception as e:
            print(f"  ❌ 修复PYTHONPATH失败: {e}")
    
    def _fix_relative_imports(self, issue):
        """修复相对导入问题"""
        try:
            # 在项目根目录创建__init__.py文件
            init_files = [
                self.project_root / "__init__.py",
                self.project_root / "modules" / "__init__.py",
                self.project_root / "apps" / "__init__.py",
                self.project_root / "config" / "__init__.py",
                self.project_root / "scripts" / "__init__.py"
            ]
            
            for init_file in init_files:
                if not init_file.exists():
                    init_file.parent.mkdir(parents=True, exist_ok=True)
                    with open(init_file, 'w', encoding='utf-8') as f:
                        f.write('# -*- coding: utf-8 -*-\n')
                    
                    self.fixes_applied.append({
                        "type": "init_file",
                        "description": f"创建 {init_file.name}",
                        "file": str(init_file)
                    })
                    
        except Exception as e:
            print(f"  ❌ 修复相对导入失败: {e}")
    
    def _fix_dependencies(self, issue):
        """修复依赖问题"""
        try:
            # 创建requirements.txt文件
            requirements_file = self.project_root / "requirements.txt"
            
            common_requirements = [
                "requests>=2.25.0",
                "psutil>=5.8.0", 
                "colorama>=0.4.4",
                "cryptography>=3.4.0",
                "paramiko>=2.7.0",
                "scapy>=2.4.0",
                "netifaces>=0.11.0"
            ]
            
            with open(requirements_file, 'w', encoding='utf-8') as f:
                f.write('\n'.join(common_requirements))
            
            self.fixes_applied.append({
                "type": "requirements",
                "description": "创建requirements.txt文件",
                "file": str(requirements_file),
                "install_command": "pip install -r requirements.txt"
            })
            
        except Exception as e:
            print(f"  ❌ 创建requirements.txt失败: {e}")
    
    def _fix_missing_modules(self, issue):
        """修复缺失模块问题"""
        try:
            # 创建简单的模块占位符
            if "modules." in issue.get("description", ""):
                modules_dir = self.project_root / "modules"
                modules_dir.mkdir(exist_ok=True)
                
                # 创建__init__.py
                init_file = modules_dir / "__init__.py"
                if not init_file.exists():
                    with open(init_file, 'w', encoding='utf-8') as f:
                        f.write('# -*- coding: utf-8 -*-\n')
                
                self.fixes_applied.append({
                    "type": "module_placeholder",
                    "description": "创建modules目录结构",
                    "file": str(modules_dir)
                })
                
        except Exception as e:
            print(f"  ❌ 修复缺失模块失败: {e}")
    
    def generate_report(self):
        """生成修复报告"""
        report = {
            "timestamp": datetime.now().isoformat(),
            "project_root": str(self.project_root),
            "issues_found": len(self.issues_found),
            "fixes_applied": len(self.fixes_applied),
            "issues": self.issues_found,
            "fixes": self.fixes_applied,
            "recommendations": self._generate_recommendations()
        }
        
        # 保存JSON报告
        report_file = self.project_root / "reports" / "execution_fix_report.json"
        report_file.parent.mkdir(exist_ok=True)
        
        with open(report_file, 'w', encoding='utf-8') as f:
            json.dump(report, f, indent=2, ensure_ascii=False)
        
        # 生成Markdown报告
        self._generate_markdown_report(report)
        
        return report
    
    def _generate_recommendations(self):
        """生成建议"""
        recommendations = []
        
        if any(issue["type"] == "relative_import" for issue in self.issues_found):
            recommendations.append("建议重构代码，使用绝对导入替代相对导入")
        
        if any(issue["type"] == "missing_dependencies" for issue in self.issues_found):
            recommendations.append("运行 'pip install -r requirements.txt' 安装依赖")
        
        if any(issue["type"] == "pythonpath" for issue in self.issues_found):
            recommendations.append("在运行脚本前先执行 'python set_environment.py'")
        
        recommendations.append("建议在虚拟环境中运行项目")
        recommendations.append("定期更新依赖包到最新版本")
        
        return recommendations
    
    def _generate_markdown_report(self, report):
        """生成Markdown格式报告"""
        report_file = self.project_root / "reports" / "execution_fix_report.md"
        
        with open(report_file, 'w', encoding='utf-8') as f:
            f.write("# 脚本执行问题修复报告\n\n")
            f.write(f"**生成时间**: {report['timestamp']}\n")
            f.write(f"**项目路径**: {report['project_root']}\n\n")
            
            f.write("## 📊 修复统计\n\n")
            f.write(f"- **发现问题**: {report['issues_found']} 个\n")
            f.write(f"- **应用修复**: {report['fixes_applied']} 个\n\n")
            
            if report['issues']:
                f.write("## 🔍 发现的问题\n\n")
                for i, issue in enumerate(report['issues'], 1):
                    severity_emoji = {"high": "🔴", "medium": "🟡", "low": "🟢"}
                    emoji = severity_emoji.get(issue['severity'], '⚪')
                    
                    f.write(f"### {i}. {emoji} {issue['description']}\n")
                    f.write(f"- **类型**: {issue['type']}\n")
                    f.write(f"- **严重程度**: {issue['severity']}\n")
                    f.write(f"- **解决方案**: {issue['solution']}\n\n")
            
            if report['fixes']:
                f.write("## 🔧 应用的修复\n\n")
                for i, fix in enumerate(report['fixes'], 1):
                    f.write(f"### {i}. {fix['description']}\n")
                    f.write(f"- **类型**: {fix['type']}\n")
                    if 'file' in fix:
                        f.write(f"- **文件**: {fix['file']}\n")
                    if 'install_command' in fix:
                        f.write(f"- **安装命令**: `{fix['install_command']}`\n")
                    f.write("\n")
            
            if report['recommendations']:
                f.write("## 💡 建议\n\n")
                for rec in report['recommendations']:
                    f.write(f"- {rec}\n")

def main():
    """主函数"""
    print("🚀 脚本执行问题修复工具启动")
    print("=" * 50)
    
    # 初始化修复器
    fixer = ExecutionIssueFixer()
    
    # 诊断问题
    issues = fixer.diagnose_issues()
    
    print(f"\n📋 诊断完成，发现 {len(issues)} 个问题")
    
    if issues:
        # 显示问题摘要
        print("\n🔍 问题摘要:")
        for issue in issues:
            severity_emoji = {"high": "🔴", "medium": "🟡", "low": "🟢"}
            emoji = severity_emoji.get(issue['severity'], '⚪')
            print(f"  {emoji} {issue['description']}")
        
        # 应用修复
        print("\n🔧 开始应用修复...")
        fixer.apply_fixes()
        
        print(f"✅ 应用了 {len(fixer.fixes_applied)} 个修复")
    else:
        print("✅ 未发现执行问题")
    
    # 生成报告
    print("\n📄 生成修复报告...")
    report = fixer.generate_report()
    
    print(f"📊 报告已保存到: reports/execution_fix_report.json")
    print(f"📝 Markdown报告: reports/execution_fix_report.md")
    
    # 显示建议
    if report['recommendations']:
        print("\n💡 建议:")
        for rec in report['recommendations']:
            print(f"  • {rec}")
    
    print("\n🎯 修复完成！")

if __name__ == "__main__":
    main()